#define DEBUG_TYPE "mpi-expand-window"
#include "llvm/Transforms/MPI/MPI.h"
#include "llvm/Analysis/MemoryDependenceAnalysis.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Pass.h"
#include "llvm/Analysis/InlineCost.h"
#include "llvm/IntrinsicInst.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/MPI/MPI.h"
#include <vector>
#include <map>
#include <functional>
#include <utility>

using namespace llvm;

// FIXME: Why should I write this again? It's in "llvm/Analysis/InlineCost.h"
// already. Not writing this cause compiling error. Why?
namespace llvm {
	bool callIsSmall(const Function *Callee);
}

namespace {
typedef BasicBlock::iterator BI;

struct alwaystrue
{
	bool operator()(BI, BI) const
	{
		return true;
	}
};

// FIXME: exclusively for MPIOverlapWindowExpansion::getOverlapWindow
// No good for encapsulation.
class RequestHandleMatch : public std::unary_function<BI,bool> {
	const BI initCall;
public:
	explicit RequestHandleMatch(const BI cs) : initCall(cs) {}
	bool operator()(const BI waitCall) const {
		// Test if the same request handle
		return initCall->getOperand(7) == waitCall->getOperand(1);
	}
};

struct CodeMetricsExt : public CodeMetrics {
	void analyzeInstRange(BI begin, BI end);
};

class MPIOverlapWindowExpansion : public FunctionPass {
private:
	AliasAnalysis *AA;
	DominatorTree *DT;
	PostDominatorTree *PDT;
	LoopInfo *LI;

public:
	static char ID;
	MPIOverlapWindowExpansion() : FunctionPass(ID) {}

	// Estimate computation time in the overlap window.
	// Use the average time of executing one instrucion as time unit.
	unsigned estComputationTime(BI begin, BI end);
	// Estimate data transfer time of a send or receive call
	// Use the average time of executing one instrucion as time unit.
	unsigned estDataTransferTime(CallInst* c);
	// bb is input argument, windowsMap is output argument
	// windowsMap is a stl map that key is send/recv call
	// value is wait call, both of BasicBlock::iterator type.
	void getOverlapWindow(BasicBlock* bb,
			std::map<BI, BI, alwaystrue> &windowsMap);
	virtual bool runOnFunction(Function &F);

	virtual void getAnalysisUsage(AnalysisUsage &AU) const {
		AU.addRequired<AliasAnalysis>();
		AU.addRequired<MemoryDependenceAnalysis>();
		AU.addRequired<DominatorTree>();
		AU.addRequired<PostDominatorTree>();
		AU.addRequired<LoopInfo>();
	};
};

}

char MPIOverlapWindowExpansion::ID = 0;
static RegisterPass<MPIOverlapWindowExpansion>
X("mpi-owexp", "MPI Overlap Window Expansion Pass");
//INITIALIZE_PASS(MPIOverlapWindowExpansion, "mpi-owexp", "MPI Overlap Window Expansion Pass", false, false);

FunctionPass *createMPIOverlapWindowExpansionPass() {
	return new MPIOverlapWindowExpansion();
}

// Copied from LoopUnrollPass
/// ApproximateLoopSize - Approximate the size of the loop.
static unsigned ApproximateLoopSize(const Loop *L, unsigned &NumCalls) {
	CodeMetrics Metrics;
	for (Loop::block_iterator I = L->block_begin(), E = L->block_end();
			I != E; ++I)
		Metrics.analyzeBasicBlock(*I);
	NumCalls = Metrics.NumCalls;
	return Metrics.NumInsts;
}

// Adapted from llvm::CodeMetrics::analyzeBasicBlock.
void CodeMetricsExt::analyzeInstRange(BI II, BI E) {
	  for (;II != E; ++II) {
	    if (isa<PHINode>(II)) continue;           // PHI nodes don't count.

	    // Special handling for calls.
	    if (isa<CallInst>(II) || isa<InvokeInst>(II)) {
	      if (isa<DbgInfoIntrinsic>(II))
	        continue;  // Debug intrinsics don't count as size.

	      CallSite CS = CallSite::get(const_cast<Instruction*>(&*II));


	      if (!isa<IntrinsicInst>(II) && !callIsSmall(CS.getCalledFunction())) {
	        // Each argument to a call takes on average one instruction to set up.
	        NumInsts += CS.arg_size();
	        ++NumCalls;
	      }
	    }

	    if (const AllocaInst *AI = dyn_cast<AllocaInst>(II)) {
	      if (!AI->isStaticAlloca())
	        this->usesDynamicAlloca = true;
	    }

	    if (isa<ExtractElementInst>(II) || II->getType()->isVectorTy())
	      ++NumVectorInsts;

	    if (const CastInst *CI = dyn_cast<CastInst>(II)) {
	      // Noop casts, including ptr <-> int,  don't count.
	      if (CI->isLosslessCast() || isa<IntToPtrInst>(CI) ||
	          isa<PtrToIntInst>(CI))
	        continue;
	      // Result of a cmp instruction is often extended (to be used by other
	      // cmp instructions, logical or return instructions). These are usually
	      // nop on most sane targets.
	      if (isa<CmpInst>(CI->getOperand(0)))
	        continue;
	    } else if (const GetElementPtrInst *GEPI = dyn_cast<GetElementPtrInst>(II)){
	      // If a GEP has all constant indices, it will probably be folded with
	      // a load/store.
	      if (GEPI->hasAllConstantIndices())
	        continue;
	    }

	    ++NumInsts;
	  }
}

unsigned MPIOverlapWindowExpansion::estComputationTime(BI begin, BI end) {
	CodeMetricsExt MetricsExt;
	MetricsExt.analyzeInstRange(begin, end);
	// TODO: Use various statistics in CodeMetricsExt to estimate Computation
	// time.
	return MetricsExt.NumInsts;
}

unsigned MPIOverlapWindowExpansion::estDataTransferTime(CallInst* C) {
	// C must be either MPI_Isend or MPI_Irecv
	Function *F = C->getCalledFunction();
	if(F->getName() != "MPI_Isend" && F->getName() != "MPI_Irecv")
			return 0;

	// TODO: analysis MPI_Isend/MPI_Irecv's transfer buffer. Develop good
	// heuristic to estimate data transfer time, potential factors include
	// history(profiling), data type size(msg size), destination(distance),
	// tag(msg type), communicator(node topology), etc. Could be arbitrary
	// complex.

	// Transferring take forever that we never get big enough window.
	return 1000;

}

void MPIOverlapWindowExpansion::getOverlapWindow(BasicBlock* bb,
								std::map<BI, BI, alwaystrue> &windowsMap) {
	std::vector<BI> initVec;
	std::vector<BI> waitVec;

	// Compute initVec and waitVec
	for(BI II= bb->begin(), IE = bb->end(); II != IE; ++II) {
		if (!isa<CallInst>(*II) && !isa<InvokeInst>(*II))
			continue;

		CallSite CS = CallSite::get(II);
		Function *F = CS.getCalledFunction();
		if(F->getName() == "MPI_Isend" || F->getName() == "MPI_Irecv") {
			initVec.push_back(II);
		}

		if(F->getName() == "MPI_Wait") {
			waitVec.push_back(II);
		}
	}

	// Testing purpose
	DEBUG(errs() << "Init calls and waits:\n\n");
	DEBUG(for_each(initVec.begin(), initVec.end(),
							std::mem_fun(&Instruction::dump)));
	DEBUG(for_each(waitVec.begin(), waitVec.end(),
							std::mem_fun(&Instruction::dump)));

	// Match init and wait
	for(std::vector<BI>::iterator VI = initVec.begin(), VE = initVec.end();
															VI != VE; VI++) {
		std::vector<BI>::iterator p =
			find_if(waitVec.begin(), waitVec.end(), RequestHandleMatch(*VI));
		if(p != waitVec.end()) {
			windowsMap[*VI] = *p;
		} else {
			DEBUG(errs() << "some init Call has no corresponding wait call\n");
			exit(1);
		}
	}

	// Testing purpose
	DEBUG(errs() << "Matching init calls and waits:\n\n");
	typedef std::map<BI, BI>::iterator CI;
	for(CI p = windowsMap.begin(); p != windowsMap.end(); p++) {
		(p->first)->dump();
		(p->second)->dump();
		errs() << "\n";
	}
}


bool MPIOverlapWindowExpansion::runOnFunction(Function &F) {
	bool Changed = false;

	AA = &getAnalysis<AliasAnalysis>();
	DT = &getAnalysis<DominatorTree>();
	PDT = &getAnalysis<PostDominatorTree>();
	LI = &getAnalysis<LoopInfo>();

	for(Function::iterator FI = F.begin(), FE = F.end(); FI != FE; ++FI) {
		BasicBlock *BB = &*FI;

		// Find matching init & wait in this Basic Block.
		std::map<BI, BI, alwaystrue> windowsMap;
		getOverlapWindow(BB, windowsMap);

		// FIXME: For now we are only doing the very simple case of bringing loops
		// inside the overlap window, and we reject the cases where we are inside
		// one
		if(LI->getLoopFor(BB))
			continue;


		for(BI II = BB->begin(), IE = BB->end(); II != IE; ++II) {
			// FIXME: Not sure if this this necessary. All the MPI_Isend &
			// MPI_Irecv is tranlated into callInst, but not InvokeInst AFAIK.
			if (!isa<CallInst>(*II) && !isa<InvokeInst>(*II))
				continue;

			CallSite CS = CallSite::get(II);
			Function *F = CS.getCalledFunction();

			// Here, we find a a communication call, bring ahead loops right away.
			// TODO: When we find a communication call, find that corresponding wait
			// immediately, then calculate if overlapped computation time > data
			// transfer time. If not, seek chance to do variable cloning & loop
			// fission.
			if(F->getName() != "MPI_Isend" && F->getName() != "MPI_Irecv")
				continue;

			// Make sure we have the right signature
			if(!MPISupport::verifyIrecvSignature(F->getFunctionType(), &F->getContext()))
				continue;

			// If overlapped computation time > data transfer time, no need to expand
			// overlap window.
			if(estComputationTime(II,windowsMap[II]) >
								estDataTransferTime(dyn_cast<CallInst>(&*II)))
				continue;

			// Use existing LLVM passes
			// TODO: Identify chance to do varible cloning
			// TODO: Identify chance to do loop fission
			// TODO: Identify chance to do CCTP (loop transformation)

			// TODO: We see if we are immediately dominated by a loop that we can move
			// otherwise we give up. This is to keep it very very simple.
			// for(pred_iterator PB = pred_begin(BB), PE = pred_end(BB); PB != PE; ++PB) {
			BasicBlock *PredBB = BB->getSinglePredecessor();
			if(!PredBB)
				continue;

			// Make sure the predecessor basic block is a loop
			Loop *L = 0;
			if(!(L = LI->getLoopFor(PredBB)))
				continue;

			// since our BB is not inside a loop, only outermost loops can be preds
			assert(!L->getParentLoop() && "Predecessor BB cannot be in inner loop.");

			if (!L->hasDedicatedExits())
				continue;

			// Deal only with loops that don't have inner loops (to make it simple)
			if (L->begin() != L->end())
				return false;

			// Make sure we can move this loop around (dom/post-dom constraints)
			BasicBlock *LoopHeader = L->getLoopPreheader();
			if(!LoopHeader)
				LoopHeader = L->getHeader();

			if(!DT->properlyDominates(LoopHeader, BB))
				continue;

			// BB must post-dominate the loop, FIXME: this is probably not enough
			if(!PDT->properlyDominates(BB, LoopHeader))
				continue;

			// Make sure we don't have any dependencies on this loop (very basic):
			// TODO: We do only alias analysis for now, to test our MPIAliasAnalysis,
			// but we also need to check that we are not putting defines after uses
			// See LoopDeletion, LICM, etc
			const Value *IBuf = MPISupport::getPointerArg(CS, MPIArgs::IBuf);
			const Value *IDatatype = MPISupport::getPointerArg(CS, MPIArgs::IDatatype);
			const Value *IComm = MPISupport::getPointerArg(CS, MPIArgs::IComm);
			const Value *IRequest = MPISupport::getPointerArg(CS, MPIArgs::IRequest);
			bool CanModify = false;
			for (Loop::block_iterator LBI = L->block_begin(), LBE = L->block_end();
					LBI != LBE && !CanModify; ++LBI) {
				BasicBlock *LoopBB = *LBI;
				// Make sure there are no instruction in the basic block that modify
				// the locations references by the pointers passed to our call,
				// otherwise moving the loop after call would change semantics.
				if(AA->canBasicBlockModify(*LoopBB, IBuf, ~0U)) {
					errs() << "\nCall to " << F->getName()
                				 << " conflicts with loop on 'buf' parameter.\n\n";
					CanModify = true;
				}

				if(AA->canBasicBlockModify(*LoopBB, IDatatype, ~0U)) {
					errs() << "\nCall to " << F->getName()
                				 << " conflicts with loop on 'datatype' parameter.\n\n";
					CanModify = true;
				}

				if(AA->canBasicBlockModify(*LoopBB, IComm, ~0U)) {
					errs() << "\nCall to " << F->getName()
                				 << " conflicts with loop on 'comm' parameter.\n\n";
					CanModify = true;
				}

				if(AA->canBasicBlockModify(*LoopBB, IRequest, ~0U)) {
					errs() << "\nCall to " << F->getName()
                				 << " conflicts with loop on 'request' parameter.\n\n";
					CanModify = true;
				}

				// now see if we have call sites that interfere (i.e.
				// hidden effects not directly shown on parameters)
				for(BasicBlock::iterator LI = LoopBB->begin(), LE = LoopBB->end();
						LI != LE && !CanModify; ++LI) {

					if (!isa<CallInst>(*LI) && !isa<InvokeInst>(*LI))
						continue;

					if(AA->getModRefInfo(CS, CallSite::get(LI)) & AliasAnalysis::Mod ||
							AA->getModRefInfo(CallSite::get(LI), CS) & AliasAnalysis::Mod) {
						CanModify = true;
						errs() << "\nCall " << F->getName()
                				   << " conflicts with loop on hidden effects.\n\n";
					}
				}
				// TODO: We also need to make sure that the instructions before this call site
				// on this BB don't use the values computed inside the loop.
			}

			if(CanModify)
				continue;

			// Insert the loop after this instruction.
			// Note that we kill all the previous passes when we do this.
			errs() << "\nFrom an AA point-of-view, the loop:\n\n";
			L->dump();
			errs() << "\nCould be moved after the instruction:\n\n";
			CS.getInstruction()->dump();
		}
	}

	return Changed;
}

